import os
import pandas as pd
from common.knowledge_rag import RAG
from common.utils import read_excel_file, save_eval_result, get_maas_models, write_to_excel, get_filename_without_ext, get_files_to_process
from common.config import config
from typing import Dict, Any, List, Optional
from ragas import EvaluationDataset
from ragas import evaluate
from ragas.metrics import (
    LLMContextPrecisionWithoutReference,
    LLMContextRecall,
    ResponseRelevancy,
    Faithfulness)
from tqdm import tqdm


def get_test_set_from_df(df: pd.DataFrame) -> List[Dict[str, str]]:
    """
    从DataFrame中获取问题、参考答案和参考上下文列表
    
    Args:
        df: pandas DataFrame，包含问题、参考答案和参考上下文列
    Returns:
        list: 包含问题、参考答案和参考上下文的字典列表
    """
    result = []
    if '问题' in df.columns and '参考答案' in df.columns:
        # 获取所有参考上下文列
        context_columns = [col for col in df.columns if col.startswith('参考上下文')]
        
        for _, row in df.iterrows():
            # 获取该行的所有非空参考上下文
            contexts = [
                str(row[col]) 
                for col in context_columns 
                if pd.notna(row[col])  # 排除NaN值
            ]
            
            result.append({
                'user_input': row['问题'],
                'reference': row['参考答案'],
                'reference_contexts': contexts
            })
    return result


def process_single_question(rag: RAG, item: Dict[str, str]) -> Dict[str, Any]:
    """
    处理单个问题，调用RAG接口并处理返回结果
    
    Args:
        rag: RAG实例
        item: 包含问题和参考答案的字典
    Returns:
        dict: 处理结果
    """
    try:
        response = rag.generate(item['user_input'])
        retrieved_contexts = [
            result['recall_text']
            for doc in response['related']['documents']
            for result in doc['retrieve_results']
        ]
        
        # 创建新的字典，保持原有的key
        result = item.copy()
        # 添加新的key
        result.update({
            'response': response.get('message', ''),
            'retrieved_contexts': retrieved_contexts
        })
        
        return result
    except Exception as e:
        print(f"处理问题失败: {str(e)}")
        return item


def collection_evaluation_data(rag: RAG, file: str) -> Optional[List[Dict[str, Any]]]:
    """
    收集评估数据
    """
    file_path = os.path.join(config.path.results_path, file)
    print(f"\n处理文件: {file}")
    
    # 读取当前文件
    df = read_excel_file(file_path)
    if df.empty:
        print(f"文件 {file} 为空")
        return
    
    # 获取问题和参考答案列表
    items = get_test_set_from_df(df)
    if not items:
        print(f"文件 {file} 中未找到问题和参考答案")
        return

    # 处理每个问题
    dataset = []
    for item in tqdm(items, desc="处理问题"):
        result = process_single_question(rag, item)
        dataset.append(result)
    return dataset


def main(metrics: List[Any], specific_files: list[str] | None = None):
    """
    评估
    
    Args:
        metrics: 评估指标列表
        specific_files: 指定要处理的文件列表，如果为None则处理results_path目录下的所有文件
    """
    try:
        # 初始化RAG实例
        rag = RAG(base_url=config.rag.base_url)
        print("RAG初始化成功，登录接口调用成功")

        # 获取需要处理的文件列表
        print(f"测试集目录: {config.path.results_path}")
        files_to_process = get_files_to_process(specific_files, config.path.results_path)
        if not files_to_process:
            return
        
        print(f"将处理以下 {len(files_to_process)} 个文件:")
        for file in files_to_process:
            print(f"  - {os.path.basename(file)}")

        # 处理文件
        for file_path in tqdm(files_to_process, desc="收集评估数据"):
            # 获取文件名（不带后缀）
            file_name_without_ext = get_filename_without_ext(file_path)
            # 收集评估数据
            dataset = collection_evaluation_data(rag, file_path)
            # 创建评估数据集
            evaluation_dataset = EvaluationDataset.from_list(dataset)
            # 配置模型
            evaluator_llm, evaluator_embeddings = get_maas_models(
                base_url=config.maas.base_url,
                embedding_model=config.maas.embedding_model,
                model=config.maas.evaluator_model
            )
            
            # 评估
            result = evaluate(
                dataset=evaluation_dataset,
                metrics=metrics,
                llm=evaluator_llm,
                embeddings=evaluator_embeddings
            )
            
            # 保存评估结果
            eval_result_file = os.path.join(config.path.eval_path, f"{file_name_without_ext}_eval_result.xlsx")
            save_eval_result(
                repr_dict=result._repr_dict,
                file_name=file_path,
                file_path=eval_result_file
            )
            # 保存每一步的评估结果
            step_result_file = os.path.join(config.path.eval_path, f"{file_name_without_ext}_eval_step_result.xlsx")
            write_to_excel(result.to_pandas(), step_result_file)
    except Exception as e:
        print(f"RAG初始化失败: {str(e)}")


if __name__ == "__main__":
    # 指定要使用的评估指标
    metrics = [
        LLMContextPrecisionWithoutReference(), # 上下文精度
        LLMContextRecall(), # 上下文召回率
        ResponseRelevancy(), # 响应相关性
        Faithfulness(), # 忠实度
    ]
    
    # # 获取所有测试集文件
    # all_files = get_files_to_process(file_path=config.path.results_path)
    # # 筛选出merged_no_duplicates结尾的文件
    # merged_files = [
    #     f for f in all_files 
    #     if f.endswith('_merged_no_duplicates.xlsx')
    # ]
    merged_files = [
        "/home/yunzhihui/zane/rag_eval/test_sets/GraphRAG：革新检索增强生成的图基技术探索_merged_no_duplicates.xlsx",
        "/home/yunzhihui/zane/rag_eval/test_sets/RAGAS：RAG评估测试集生成技术的全面剖析与应用研究_merged_no_duplicates.xlsx",
    ]
    main(metrics=metrics, specific_files=merged_files)
